msg_tool\scripts\kirikiri/
tjs_ns0.rs

1//! Kirikiri TJS NS0 binary encoded script
2use crate::ext::io::*;
3use crate::scripts::base::*;
4use crate::types::*;
5use crate::utils::encoding::{decode_to_string, encode_string};
6use crate::utils::struct_pack::*;
7use anyhow::Result;
8use msg_tool_macro::*;
9use overf::wrapping;
10use serde::{Deserialize, Serialize};
11use std::collections::BTreeMap;
12use std::io::{Read, Seek, Write};
13
14#[derive(Debug)]
15/// Kirikiri TJS NS0 Script Builder
16pub struct TjsNs0Builder {}
17
18impl TjsNs0Builder {
19    /// Creates a new instance of `TjsNs0Builder`
20    pub fn new() -> Self {
21        Self {}
22    }
23}
24
25impl ScriptBuilder for TjsNs0Builder {
26    fn default_encoding(&self) -> Encoding {
27        Encoding::Utf16LE
28    }
29
30    fn build_script(
31        &self,
32        buf: Vec<u8>,
33        filename: &str,
34        encoding: Encoding,
35        _archive_encoding: Encoding,
36        config: &ExtraConfig,
37        _archive: Option<&Box<dyn Script>>,
38    ) -> Result<Box<dyn Script>> {
39        Ok(Box::new(TjsNs0::new(buf, filename, encoding, config)?))
40    }
41
42    fn extensions(&self) -> &'static [&'static str] {
43        &["pbd", "tjs"]
44    }
45
46    fn script_type(&self) -> &'static ScriptType {
47        &ScriptType::KirikiriTjsNs0
48    }
49
50    fn is_this_format(&self, _filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
51        if buf_len >= 8 && (buf.starts_with(b"TJS/ns0\0") || buf.starts_with(b"TJS/4s0\0")) {
52            return Some(100);
53        }
54        None
55    }
56
57    fn can_create_file(&self) -> bool {
58        true
59    }
60
61    fn create_file<'a>(
62        &'a self,
63        filename: &'a str,
64        mut writer: Box<dyn WriteSeek + 'a>,
65        encoding: Encoding,
66        file_encoding: Encoding,
67        config: &ExtraConfig,
68    ) -> Result<()> {
69        let s = crate::utils::files::read_file(filename)?;
70        let s = decode_to_string(file_encoding, &s, true)?;
71        let data: TjsValue = if config.custom_yaml {
72            serde_yaml_ng::from_str(&s)?
73        } else {
74            serde_json::from_str(&s)?
75        };
76        let header = Header {
77            magic: *b"TJS/",
78            check: *b"ns0\0",
79            seed: u32::from_le_bytes(*b"TJS\0"),
80            crypt: 0,
81            iv_len: 0,
82        };
83        let mut checker = ByteChecker::new(header.seed);
84        header.pack(&mut writer, false, encoding)?;
85        data.pack(&mut checker, &mut writer, false, encoding)?;
86        let checksum = checker.final_check();
87        writer.write_u32(checksum)?;
88        writer.flush()?;
89        Ok(())
90    }
91}
92
93#[derive(Debug, Serialize, Deserialize)]
94#[serde(untagged)]
95enum TjsValue {
96    Void(()),
97    Int(i64),
98    Double(f64),
99    Str(String),
100    Array(Vec<TjsValue>),
101    Dict(BTreeMap<String, TjsValue>),
102}
103
104fn unpack_string<R: Read + Seek>(reader: &mut R, big: bool, encoding: Encoding) -> Result<String> {
105    let len = u32::unpack(reader, big, encoding)? as usize;
106    let tlen = if encoding.is_utf16le() { len * 2 } else { len };
107    let mut buf = vec![0u8; tlen];
108    reader.read_exact(&mut buf)?;
109    let s = decode_to_string(encoding, &buf, true)?;
110    Ok(s)
111}
112
113fn pack_string<W: Write>(s: &str, writer: &mut W, big: bool, encoding: Encoding) -> Result<()> {
114    let encoded = encode_string(encoding, s, false)?;
115    let len = if encoding.is_utf16le() {
116        (encoded.len() / 2) as u32
117    } else {
118        encoded.len() as u32
119    };
120    len.pack(writer, big, encoding)?;
121    writer.write_all(&encoded)?;
122    Ok(())
123}
124
125impl TjsValue {
126    fn pack<W: Write>(
127        &self,
128        checker: &mut ByteChecker,
129        writer: &mut W,
130        big: bool,
131        encoding: Encoding,
132    ) -> Result<()> {
133        match self {
134            Self::Void(()) => {
135                let typ_byte = 0;
136                let check_byte = checker.get_seed(typ_byte);
137                let typ = ((check_byte as u16) << 8) | (typ_byte as u16);
138                typ.pack(writer, big, encoding)?;
139            }
140            Self::Str(s) => {
141                let typ_byte = 2;
142                let check_byte = checker.get_seed(typ_byte);
143                let typ = ((check_byte as u16) << 8) | (typ_byte as u16);
144                typ.pack(writer, big, encoding)?;
145                pack_string(s, writer, big, encoding)?;
146            }
147            Self::Int(i) => {
148                let typ_byte = 4;
149                let check_byte = checker.get_seed(typ_byte);
150                let typ = ((check_byte as u16) << 8) | (typ_byte as u16);
151                typ.pack(writer, big, encoding)?;
152                i.pack(writer, big, encoding)?;
153            }
154            Self::Double(f) => {
155                let typ_byte = 5;
156                let check_byte = checker.get_seed(typ_byte);
157                let typ = ((check_byte as u16) << 8) | (typ_byte as u16);
158                typ.pack(writer, big, encoding)?;
159                f.pack(writer, big, encoding)?;
160            }
161            Self::Array(arr) => {
162                let typ_byte = 0x81;
163                let check_byte = checker.get_seed(typ_byte);
164                let typ = ((check_byte as u16) << 8) | (typ_byte as u16);
165                typ.pack(writer, big, encoding)?;
166                let arr_len = arr.len() as u32;
167                arr_len.pack(writer, big, encoding)?;
168                for item in arr {
169                    item.pack(checker, writer, big, encoding)?;
170                }
171            }
172            Self::Dict(dict) => {
173                let typ_byte = 0xC1;
174                let check_byte = checker.get_seed(typ_byte);
175                let typ = ((check_byte as u16) << 8) | (typ_byte as u16);
176                typ.pack(writer, big, encoding)?;
177                let dict_len = dict.len() as u32;
178                dict_len.pack(writer, big, encoding)?;
179                for (key, value) in dict {
180                    pack_string(key, writer, big, encoding)?;
181                    value.pack(checker, writer, big, encoding)?;
182                }
183            }
184        }
185        Ok(())
186    }
187
188    fn unpack<R: Read + Seek>(
189        checker: &mut ByteChecker,
190        reader: &mut R,
191        big: bool,
192        encoding: Encoding,
193    ) -> Result<Self> {
194        let typ = u16::unpack(reader, big, encoding)?;
195        let typ_byte = (typ & 0xff) as u8;
196        let check_byte = (typ >> 8) as u8;
197        let expected_check = checker.get_seed(typ_byte);
198        if check_byte != expected_check {
199            return Err(anyhow::anyhow!(
200                "TJS/ns0 byte check failed: expected {}, got {} at pos {}",
201                expected_check,
202                check_byte,
203                reader.stream_position()? - 1
204            ));
205        }
206        Ok(match typ_byte {
207            0 => TjsValue::Void(()),
208            2 => TjsValue::Str(unpack_string(reader, big, encoding)?),
209            4 => TjsValue::Int(i64::unpack(reader, big, encoding)?),
210            5 => TjsValue::Double(f64::unpack(reader, big, encoding)?),
211            0x81 => {
212                let arr_len = u32::unpack(reader, big, encoding)? as usize;
213                let mut arr = Vec::with_capacity(arr_len);
214                for _ in 0..arr_len {
215                    arr.push(TjsValue::unpack(checker, reader, big, encoding)?);
216                }
217                TjsValue::Array(arr)
218            }
219            0xC1 => {
220                let kv_len = u32::unpack(reader, big, encoding)? as usize;
221                let mut dict = BTreeMap::new();
222                for _ in 0..kv_len {
223                    let key = unpack_string(reader, big, encoding)?;
224                    let value = TjsValue::unpack(checker, reader, big, encoding)?;
225                    dict.insert(key, value);
226                }
227                TjsValue::Dict(dict)
228            }
229            _ => {
230                return Err(anyhow::anyhow!(
231                    "Unsupported TJS/ns0 value type: {} at pos {}",
232                    typ_byte,
233                    reader.stream_position()? - 2
234                ));
235            }
236        })
237    }
238}
239
240#[derive(Debug)]
241/// Kirikiri TJS NS0 Script
242pub struct TjsNs0 {
243    data: TjsValue,
244    custom_yaml: bool,
245    header: Header,
246}
247
248struct ByteChecker {
249    seed: u32,
250}
251
252impl ByteChecker {
253    pub fn new(seed: u32) -> Self {
254        Self { seed }
255    }
256
257    fn calculate_round(seed: &mut [u8; 4]) {
258        let a = seed[0] ^ wrapping!(seed[0] * 2);
259        let mut b = a;
260        wrapping! {
261        b >>= 2;
262        b ^= seed[2];
263        b >>= 3;
264        b ^= seed[2];
265        b ^= a;
266        }
267
268        seed[0] = seed[1];
269        seed[1] = seed[2];
270        seed[2] = b;
271    }
272
273    pub fn get_seed(&mut self, type_code: u8) -> u8 {
274        let mut s = self.seed.to_le_bytes();
275        if type_code == 0 {
276            return s[2];
277        }
278        Self::calculate_round(&mut s);
279        self.seed = u32::from_le_bytes(s);
280        return s[2];
281    }
282
283    pub fn final_check(&mut self) -> u32 {
284        let mut s = self.seed.to_le_bytes();
285        Self::calculate_round(&mut s);
286        Self::calculate_round(&mut s);
287        Self::calculate_round(&mut s);
288        let tmp = s[0];
289        s[0] = s[2];
290        s[2] = tmp;
291        u32::from_le_bytes(s)
292    }
293}
294
295#[derive(Clone, Debug, StructPack, StructUnpack)]
296struct Header {
297    magic: [u8; 4],
298    check: [u8; 4],
299    seed: u32,
300    crypt: u16,
301    iv_len: u16,
302}
303
304impl TjsNs0 {
305    /// Creates a new `TjsNs0` script from the given buffer and filename
306    ///
307    /// * `buf` - The buffer containing the TJS/ns0 data
308    /// * `filename` - The name of the file
309    /// * `encoding` - The encoding to use for strings
310    /// * `config` - Extra configuration options
311    pub fn new(
312        buf: Vec<u8>,
313        _filename: &str,
314        encoding: Encoding,
315        config: &ExtraConfig,
316    ) -> Result<Self> {
317        let mut reader = MemReader::new(buf);
318        let header = Header::unpack(&mut reader, false, encoding)?;
319        if &header.magic != b"TJS/" {
320            return Err(anyhow::anyhow!("Not a valid TJS/ns0 file"));
321        }
322        if header.check[1] != b's' || header.check[2] != b'0' || header.check[3] != 0 {
323            return Err(anyhow::anyhow!("Not a valid TJS/ns0 file"));
324        }
325        if header.crypt != 0 {
326            return Err(anyhow::anyhow!("Encrypted TJS/ns0 files are not supported"));
327        }
328        if header.iv_len != 0 {
329            return Err(anyhow::anyhow!("TJS/ns0 files with IV are not supported"));
330        }
331        let mut reader = match header.check[0] {
332            b'n' => reader,
333            b'4' => {
334                let decompressed = lz4::block::decompress(&reader.data[reader.pos..], None)?;
335                MemReader::new(decompressed)
336            }
337            _ => {
338                return Err(anyhow::anyhow!(
339                    "Unsupported compression method in TJS/ns0 file"
340                ));
341            }
342        };
343        let mut checker = ByteChecker::new(header.seed);
344        let data = TjsValue::unpack(&mut checker, &mut reader, false, encoding)?;
345        let expected_checksum = checker.final_check();
346        let actual_checksum = reader.read_u32()?;
347        if expected_checksum != actual_checksum {
348            return Err(anyhow::anyhow!(
349                "TJS/ns0 checksum mismatch: expected {:08X}, got {:08X}",
350                expected_checksum,
351                actual_checksum
352            ));
353        }
354        Ok(Self {
355            data,
356            custom_yaml: config.custom_yaml,
357            header,
358        })
359    }
360}
361
362impl Script for TjsNs0 {
363    fn default_output_script_type(&self) -> OutputScriptType {
364        OutputScriptType::Custom
365    }
366
367    fn default_format_type(&self) -> FormatOptions {
368        FormatOptions::None
369    }
370
371    fn is_output_supported(&self, output: OutputScriptType) -> bool {
372        matches!(output, OutputScriptType::Custom)
373    }
374
375    fn custom_output_extension<'a>(&'a self) -> &'a str {
376        if self.custom_yaml { "yaml" } else { "json" }
377    }
378
379    fn custom_export(&self, filename: &std::path::Path, encoding: Encoding) -> Result<()> {
380        let s = if self.custom_yaml {
381            serde_yaml_ng::to_string(&self.data)?
382        } else {
383            serde_json::to_string_pretty(&self.data)?
384        };
385        let s = encode_string(encoding, &s, false)?;
386        let mut writer = crate::utils::files::write_file(filename)?;
387        writer.write_all(&s)?;
388        Ok(())
389    }
390
391    fn custom_import<'a>(
392        &'a self,
393        custom_filename: &'a str,
394        mut file: Box<dyn WriteSeek + 'a>,
395        encoding: Encoding,
396        output_encoding: Encoding,
397    ) -> Result<()> {
398        let s = crate::utils::files::read_file(custom_filename)?;
399        let s = decode_to_string(output_encoding, &s, true)?;
400        let data: TjsValue = if self.custom_yaml {
401            serde_yaml_ng::from_str(&s)?
402        } else {
403            serde_json::from_str(&s)?
404        };
405        let mut header = self.header.clone();
406        header.check = *b"ns0\0";
407        let mut checker = ByteChecker::new(header.seed);
408        header.pack(&mut file, false, encoding)?;
409        data.pack(&mut checker, &mut file, false, encoding)?;
410        let checksum = checker.final_check();
411        file.write_u32(checksum)?;
412        file.flush()?;
413        Ok(())
414    }
415}
msg_tool\scripts\kirikiri/tjs_ns0.rs

msg_tool\scripts\kirikiri/
tjs_ns0.rs
